#pragma once

#include "tree.hpp"

enum class UnbalanceType { LA, RA, LR, RL };  // L*/R*/LR/RL

template <typename K, typename V>
class AVL : public Tree<K, V> {
  using Tree<K, V>::is_nil;

 private:
  AVL<K, V>* get_p() const { return dynamic_cast<AVL<K, V>*>(this->p); }
  AVL<K, V>* get_l() const { return dynamic_cast<AVL<K, V>*>(this->l); }
  AVL<K, V>* get_r() const { return dynamic_cast<AVL<K, V>*>(this->r); }

  AVL<K, V>* left_rotate() override {
    return dynamic_cast<AVL<K, V>*>(Tree<K, V>::left_rotate());
  }
  AVL<K, V>* right_rotate() override {
    return dynamic_cast<AVL<K, V>*>(Tree<K, V>::right_rotate());
  }

  mutable int height{0};

  // 更新以this为根的子树的高
  // 要求子树高正确
  // 如果高改变返回true反之false
  bool update_height() const {
    int new_height = -1;
    if (is_nil) {
      new_height = 0;
    } else {
      new_height = get_l()->height > get_r()->height ? get_l()->height + 1
                                                     : get_r()->height + 1;
    }
    assert(height != -1);
    if (new_height != height) {
      height = new_height;
      return true;
    } else {
      return false;
    }
  }
  // 判断当前子树左右是否平衡
  //  要求左右子树高正确
  bool balanced() const {
    if (is_nil) return true;
    int factor = get_l()->height - get_r()->height;
    return factor <= 1 && factor >= -1;
  }
  // 找到更改节点最近的不平衡祖先
  // 同时更新高
  // 找到则返回其指针否则返回nullptr
  AVL<K, V>* closest_unbalanced() {
    AVL<K, V>* to_update = this;
    if (!to_update->balanced()) {
      return to_update;
    }
    while (to_update->update_height() && to_update->get_p()) {
      to_update = to_update->get_p();
      if (!to_update->balanced()) {
        return to_update;
      }
    }
    return nullptr;
  }
  // 对不平衡情况进行分类
  // 由最近不平衡祖先调用
  UnbalanceType classify_ub_condition() {
    if (get_l()->height > get_r()->height) {
      if (get_l()->get_l()->height >= get_l()->get_r()->height) {
        return UnbalanceType::LA;
      } else {
        return UnbalanceType::LR;
      }
    } else {
      if (get_r()->get_r()->height >= get_r()->get_l()->height) {
        return UnbalanceType::RA;
      } else {
        return UnbalanceType::RL;
      }
    }
  }

  AVL<K, V>* post() {
    AVL<K, V>* closest_ub_ancestor = closest_unbalanced();
    if (closest_ub_ancestor) {
      UnbalanceType t = closest_ub_ancestor->classify_ub_condition();
      AVL<K, V>* new_local_root = nullptr;
      switch (t) {
        case UnbalanceType::LA: {
          new_local_root = closest_ub_ancestor->right_rotate();
          new_local_root->get_r()->update_height();
          new_local_root->update_height();
          break;
        }
        case UnbalanceType::LR: {
          AVL<K, V>* new_l = closest_ub_ancestor->get_l()->left_rotate();
          new_l->get_l()->update_height();
          new_l->update_height();
          new_local_root = closest_ub_ancestor->right_rotate();
          new_local_root->get_r()->update_height();
          new_local_root->update_height();
          break;
        }
        case UnbalanceType::RA: {
          new_local_root = closest_ub_ancestor->left_rotate();
          new_local_root->get_l()->update_height();
          new_local_root->update_height();
          break;
        }
        case UnbalanceType::RL: {
          AVL<K, V>* new_r = closest_ub_ancestor->get_r()->right_rotate();
          new_r->get_r()->update_height();
          new_r->update_height();
          new_local_root = closest_ub_ancestor->left_rotate();
          new_local_root->get_l()->update_height();
          new_local_root->update_height();
          break;
        }

        default:
          break;
      }

      return new_local_root->get_p();
    } else {
      return nullptr;
    }
  }

 protected:
  Tree<K, V>* new_nil() override { return new AVL(); }
  void post_insert() override { post(); }
  void post_remove() override {
    AVL<K, V>* to_adjust = this;
    while (to_adjust) {
      to_adjust = to_adjust->post();
    }
  }
  bool extra_check() const override {
    if (is_nil) {
      return true;
    } else {
      bool l_ok = get_l()->extra_check();
      bool r_ok = get_r()->extra_check();
      bool h_ok = !update_height();
      bool b_ok = balanced();
      if (!h_ok) std::cout << "高度未正确更新" << std::endl;
      if (!b_ok) std::cout << "平衡树性质未保持" << std::endl;
      return h_ok && b_ok && l_ok && r_ok;
    }
  }
  const char* extra_info() const override { return ""; }
};